
#导入数据增强工具
import Augmentor
 
AugParas={
'prePath':'D:\wuxixian\HYM\demo\\rawImage\out\\raw',
'RawPath':'D:\wuxixian\HYM\demo\\rawImage\out\\raw',
'GtPath':'D:\wuxixian\HYM\demo\\rawImage\out\gt',
'count':40,
#图像旋转： 按照概率0.8执行，最大左旋角度10，最大右旋角度10
'rotate':{
    'probability':0.8, 
    'max_left_rotation':10, 
    'max_right_rotation':10
},
#图像左右互换： 按照概率0.5执行
'flip_left_right':{
    'probability':0.8
},
#按概率随机上下翻转
'flip_top_bottom':{
    'probability':0.5
},
#图像放大缩小： 按照概率0.8执行，面积为原始图0.85倍
'zoom_random':{
    'probability':0.3, 
    'percentage_area':0.85
},
# 弹性扭曲，类似区域扭曲的感觉
'random_distortion':{
    'probability':0.8,
    'grid_width':10,
    'grid_height':10, 
    'magnitude':20
},
# 随机区域擦除
'random_erasing':{
    'probability':0.8,
    'rectangle_area':0.3
},
# 透视形变-垂直方向形变：magnitude取（0,1），指的是形变程度
'skew_tilt':{
    'probability':0.7,
    'magnitude':0.5
},
# 透视形变-斜四角形变：magnitude取（0,1），指的是形变程度
'skew_corner':{
    'probability':0.7,
    'magnitude':1
},
# 错切变换
'shear':{
    'probability':1,
    'max_shear_left':15,
    'max_shear_right':15
}
}

def augmentor():
    global AugParas
    RawPath=AugParas['RawPath']
    GtPath=AugParas['GtPath']
    #确定原始图像存储路径以及掩码文件存储路径
    p = Augmentor.Pipeline(RawPath)
    # 是标签图的路径
    p.ground_truth(GtPath)
    
    #图像旋转： 按照概率0.8执行，最大左旋角度10，最大右旋角度10
    p.rotate(probability=AugParas['rotate']['probability'], max_left_rotation=AugParas['rotate']['max_right_rotation'], max_right_rotation=AugParas['rotate']['max_right_rotation'])
    
    #图像左右互换： 按照概率0.5执行
    p.flip_left_right(probability=AugParas['flip_left_right']['probability'])
    
    
    #按概率随机上下翻转
    p.flip_top_bottom(probability=AugParas['flip_top_bottom']['probability'])
    
    #图像放大缩小： 按照概率0.8执行，面积为原始图0.85倍
    p.zoom_random(probability=AugParas['zoom_random']['probability'], percentage_area=AugParas['zoom_random']['percentage_area'])

    
    # 弹性扭曲，类似区域扭曲的感觉
    p.random_distortion(probability=AugParas['random_distortion']['probability'],grid_width=AugParas['random_distortion']['grid_width'],
                        grid_height=AugParas['random_distortion']['grid_height'], magnitude=AugParas['random_distortion']['magnitude'])#小块变形
    
    # 随机区域擦除
    p.random_erasing(probability=AugParas['random_erasing']['probability'], rectangle_area=AugParas['random_erasing']['rectangle_area'])
    
    # 透视形变-垂直方向形变：magnitude取（0,1），指的是形变程度
    p.skew_tilt(probability=AugParas['skew_tilt']['probability'], magnitude=AugParas['skew_tilt']['magnitude'])
    
    # 透视形变-斜四角形变：magnitude取（0,1），指的是形变程度
    p.skew_corner(probability=AugParas['skew_corner']['probability'],magnitude=AugParas['skew_corner']['magnitude'])
    
    # 错切变换
    p.shear(probability=AugParas['shear']['probability'],max_shear_left=AugParas['shear']['max_shear_left'],max_shear_right=AugParas['shear']['max_shear_right'])
    
    #最终扩充的数据样本数
    p.sample(AugParas['count'])
    
    # 需要把该文件复制到上一层
    import shutil
    # 先清空
    curPath=AugParas['RawPath']+'\output\\'
    nextPath=AugParas['prePath']+'\output\\'
    try:
        import shutil
        shutil.rmtree(nextPath)
    except Exception  as e :
        print(e)
    shutil.copytree(curPath,nextPath)
    try:
        import shutil
        shutil.rmtree(curPath)
    except Exception  as e :
        print(e)
    return nextPath
